/*
 * Copyright 2016-2036 the original author or authors.
 *
 * COMMERCIAL USE OF THIS SOFTWARE WITHOUT WARRANTY IS NOT ALLOWED.
 * Use is subject to license terms! You can distribute a copy of this software
 * to others for free. This software is a non-profit and open-source project.
 * Any contribution to this project will make it better.
 * All rights reserved! Owned by Stephen Liu.
 *
 */
package com.github.arrow.function;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.EnumSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;

/**
 * @author ste7en.liu@gmail.com
 * @since 2016/7/16
 */
public class GroupingWithKeys<T, K>
        implements Collector<T, Map<K, List<T>>, Map<K, List<T>>> {

    private List<K> keys;
    private Function<T, K> keyGen;

    @SafeVarargs
    public GroupingWithKeys(Function<T, K> keyGen, K... keys) {
        if (keys != null) {
            this.keys = Arrays.asList(keys);
        }
        this.keyGen = keyGen;
    }

    @Override
    public Supplier<Map<K, List<T>>> supplier() {
        return () -> {
            Map<K, List<T>> map = new LinkedHashMap<>();
            keys.forEach(x -> map.put(x, new ArrayList<>()));
            return map;
        };
    }

    @Override
    public BiConsumer<Map<K, List<T>>, T> accumulator() {
        return (m, t) -> {
            List<T> list = m.getOrDefault(keyGen.apply(t), new ArrayList<>());
            list.add(t);
            m.put(keyGen.apply(t), list);
        };
    }

    @Override
    public BinaryOperator<Map<K, List<T>>> combiner() {
        return (m1, m2) -> {
            m1.putAll(m2);
            return m1;
        };
    }

    @Override
    public Function<Map<K, List<T>>, Map<K, List<T>>> finisher() {
        return Function.identity();
    }

    @Override
    public Set<Characteristics> characteristics() {
        return Collections.unmodifiableSet(
                EnumSet.of(Characteristics.IDENTITY_FINISH, Characteristics.CONCURRENT));
    }
}
